import dotenv
import os
import weaviate
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.language_models import BaseLanguageModel
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate,FewShotChatMessagePromptTemplate
from langchain_core.retrievers import BaseRetriever
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_weaviate import WeaviateVectorStore
from weaviate.auth import AuthApiKey
dotenv.load_dotenv()

class StepBackRetriever(BaseRetriever):
    retriever:BaseRetriever
    llm:BaseLanguageModel

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> list[Document]:

        examples = [
            {"input": "博瑞AI上有关于AI应用开发的课程吗？", "output": "博瑞有哪些AI的课程"},
            {"input": "博瑞出生在哪个和国家？", "output": "博瑞的个人经历是怎么样的"},
            {"input": "司机可以开快车吗？", "output": "司机可以做什么"}
        ]

        example_prompt = ChatPromptTemplate.from_messages([
            ("human","{input}"),
            ("ai","{output}")
        ])

        few_shot_prompt = FewShotChatMessagePromptTemplate(
            examples = examples,
            example_prompt = example_prompt
        )

        # 构建生成会问题提示
        system_prompt = "你是一个知识专家，你的任务是回退问题，将问题改述更一般的前置问题，这样更容易回答，请严格按照格式回答，不要废话"
        prompt = ChatPromptTemplate.from_messages([
            ("system",system_prompt),
            few_shot_prompt,
            ("human","{question}")
        ])
#         构建生成回退问题链
        chain = (
            {"question":RunnablePassthrough()}
            |prompt
            |self.llm
            |StrOutputParser()
            |self.retriever
        )
        return chain.invoke(query)
# 导入数据
client = weaviate.connect_to_weaviate_cloud(
    skip_init_checks = True,
    cluster_url = os.getenv("WAEVIATE_URL"),
    auth_credentials = AuthApiKey(os.getenv("WEAVIATE_KEY"))
)
embedding = OpenAIEmbeddings(model="text-embedding-3-small")
db = WeaviateVectorStore(client=client,
                         index_name="DataSetTest",
                         text_key="text",
                         embedding=embedding

)

# 创建回答回退检索器
step_back_retriver = StepBackRetriever(
    retriever = db.as_retriever(),
    llm = ChatOpenAI(model="gpt-4o-mini",temperature=0)
)

# 检索文档
docs = step_back_retriver.invoke("人工智能会让世界发生翻天覆地的变化吗")

for doc in docs:
    print(doc)